# coding:utf-8
# author: 湖南涉外经济学院信息与机电工程学院
# date: 2022.09.09

import os
import base64
import time
from flask import Blueprint, render_template, request,redirect
import numpy as np
import cv2

from PIL import Image
from koutu import upload
from koutu import do_zip_compress
index_zhengjianzhao = Blueprint("zhengjianzhao", __name__)
# 设置允许的文件格式
ALLOWED_EXTENSIONS = set(['png', 'jpg', 'bmp', 'jpeg'])
# 当前文件所在路径
basepath = os.path.dirname(__file__)


def cv2_to_base64(image):
    data = cv2.imencode('.jpg', image)[1]
    return base64.b64encode(data.tobytes()).decode('utf8')


def base64_to_cv2(b64str):
    data = base64.b64decode(b64str.encode('utf8'))
    data = np.frombuffer(data, np.uint8)
    data = cv2.imdecode(data, cv2.IMREAD_COLOR)
    return data


def allowed_file(filename):
    filename = filename.lower()
    return '.' in filename and filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS


# 更换背景算法A
def cut_person(f, background_color, save_path, thresh):
    if isinstance(background_color, str):
        if background_color == '1':
            color = [255, 0, 0]
        elif background_color == '2':
            color = [67, 142, 219]
        elif background_color == '3':
            color = [255, 255, 255]
        else:
            raise Exception('背景色设置有误')
    elif isinstance(background_color, list) or isinstance(background_color, tuple):
        color = [background_color[0], background_color[1], background_color[2], 1]
    else:
        raise Exception('背景色设置有误')
    filename = f.filename
    new_img_filename = 'color' + filename
    new_img_path = os.path.join(basepath, 'static/images/target', new_img_filename)
    if save_path == None:
        base_img_filename = os.path.join(basepath, 'static/images/target', filename)
    else:
        base_img_filename = os.path.join(basepath, 'static/images/target', filename.split("/")[0],
                                         filename.split("/")[1])
        new_img_path = os.path.join(basepath, 'static/images/target', filename.split("/")[0],
                                    'color' + filename.split("/")[1])
    # cv2.imencode('.png', rgba)[1].tofile(abs_filename)
    img = np.array(Image.open(base_img_filename)).transpose((2, 0, 1))
    person = []
    for i in range(3):
        a = img[i]
        mask = np.array((img[3] < thresh), dtype=np.uint8)
        # mask=cv2.erode(mask,None,iterations=1)
        mask = cv2.dilate(mask, None, iterations=2)
        mask = np.array(mask, dtype=bool)
        a[mask] = color[i]
        person.append(a)
    print('step10 修改背景完成')
    img = np.array(person).transpose((1, 2, 0))
    im = Image.fromarray(img)
    im.save(new_img_path)
    os.remove(base_img_filename)
    return new_img_filename


# 更换背景算法B
def change_back_groundcolor(filename, background_color):
    if isinstance(background_color, str):
        if background_color == '1':
            color = [255, 0, 0, 1]
        elif background_color == '2':
            color = [67, 142, 219, 1]
        elif background_color == '3':
            color = [255, 255, 255, 1]
        else:
            raise Exception('背景色设置有误')

    elif isinstance(background_color, list) or isinstance(background_color, tuple):
        color = [background_color[0], background_color[1], background_color[2], 1]
    else:
        raise Exception('背景色设置有误')
    base_img_filename = os.path.join(basepath, 'static/images/target', filename)
    new_img_filename = 'color' + filename
    new_img_path = os.path.join(basepath, 'static/images/target', new_img_filename)
    base_img = Image.open(base_img_filename)
    img = np.array(base_img)
    for i in range(0, img.shape[0]):
        for j in range(0, img.shape[1]):
            if img[i][j][3] < 1:
                img[i][j] = color
    print('step10 修改背景完成')
    im = Image.fromarray(img)
    im = im.convert('RGB')
    im.save(new_img_path)
    return new_img_filename


# 上传并批量抠图
@index_zhengjianzhao.route('/zhengjianzhao_p', methods=['POST', 'GET'])  # 添加路由
def zhengjianzhao_p():
    if request.method == 'POST':
        try:
            time_start = time.time()
            f = request.files.getlist('file')
            save_path2 = f[0].filename.split("/")[0]
            save_path = os.path.join('static/images/target', save_path2)
            save_path1 = save_path
            selected_color = request.form.get('selected_color')
            selected_algorithm = request.form.get('selected_algorithm')
            selected_thresh = request.form.get('describe')
            print('step1 算法设置:', selected_algorithm)
            print('step2 颜色设置:', selected_color)
            print('step3 羽化设置:', selected_thresh)
            if (selected_thresh == ''):
                selected_thresh = 100
            else:
                selected_thresh = int(selected_thresh)
            if os.path.exists(save_path):
                save_path = 'save_path'
            else:
                os.mkdir(save_path)
            for i in f:
                upload(i, save_path)
                if selected_algorithm == 'A':
                    change_color_imgfile = cut_person(i, selected_color, save_path, thresh=selected_thresh)
                else:
                    return render_template('zhengjianzhao.html')
                    # change_color_imgfile = change_back_groundcolor(filename, selected_color)
            filename = os.path.join('static/images/target', change_color_imgfile)
            print('step11 证件照已生成,保存在:', filename)
            time_end = time.time()
            time_c = time_end - time_start
            print('step12 所用时间: %.4f' % time_c, 's')
            do_zip_compress(save_path1)
            return redirect('/download/' + save_path2)
        except Exception as e:
            print(e)
            return render_template('404.html')
    return render_template('zhengjianzhao.html')

# 上传并抠图
@index_zhengjianzhao.route('/zhengjianzhao', methods=['POST', 'GET'])  # 添加路由
def zhengjianzhao():
    if request.method == 'POST':
        try:
            time_start = time.time()
            f = request.files.get('file2')
            save_path = None
            selected_color = request.form.get('selected_color')
            selected_algorithm = request.form.get('selected_algorithm')
            selected_thresh = request.form.get('describe')
            print('step1 算法设置:', selected_algorithm)
            print('step2 颜色设置:', selected_color)
            print('step3 羽化设置:', selected_thresh)
            if (selected_thresh == ''):
                selected_thresh = 100
            else:
                selected_thresh = int(selected_thresh)
            sourcefile, filename, upload_path = upload(f, save_path)
            if selected_algorithm == 'A':
                change_color_imgfile = cut_person(f, selected_color, save_path, thresh=selected_thresh)
                filename = os.path.join('static/images/target', change_color_imgfile)
            else:
                return '算法B淘汰了'
            print('step11 证件照已生成,保存在:', filename)
            time_end = time.time()
            time_c = time_end - time_start
            print('step12 所用时间: %.4f' % time_c, 's')
            return render_template('zhengjianzhao.html', sourcefile=sourcefile, filename=filename)
        except Exception as e:
            print(e)
            return render_template('404.html')
    return render_template('zhengjianzhao.html')
